{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7yuytuIllsv1"
      },
      "source": [
        "# Trax Layers\n",
        "\n",
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "height": 119
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 88481,
          "status": "ok",
          "timestamp": 1563927238895,
          "user": {
            "displayName": "Lukasz Kaiser",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
            "userId": "13267693649565518272"
          },
          "user_tz": 420
        },
        "id": "oILRLCWN_16u",
        "outputId": "3f750014-c633-4162-ad07-f3c56c273304"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "T2T: skipped importing 1 data_generators modules. OK if no other errors. Depend on _heavy or problem-specific py_binary targets if trying to use a module that was skipped.\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "matplotlib.use() must be called *before* pylab, matplotlib.pyplot,\n",
            "or matplotlib.backends is imported for the first time.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "#@title Licence and python imports.\n",
        "# Copyright 2018 Google LLC.\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "import datetime\n",
        "import numpy as onp\n",
        "\n",
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "cellView": "both",
        "colab": {
          "height": 51
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 2992,
          "status": "ok",
          "timestamp": 1563927313403,
          "user": {
            "displayName": "Lukasz Kaiser",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
            "userId": "13267693649565518272"
          },
          "user_tz": 420
        },
        "id": "vlGjGoGMTt-D",
        "outputId": "7a7b5a1e-c01e-4a5e-eeb1-a88d9500aad4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/bin/sh: pip: command not found\n",
            "/bin/sh: pip: command not found\n"
          ]
        }
      ],
      "source": [
        "#@title Install TensorFlow, Tensor2Tensor and Trax.\n",
        "\n",
        "! pip install -q tensorflow\n",
        "! pip install -q -U tensor2tensor\n",
        "\n",
        "from tensor2tensor.trax import trax\n",
        "from tensor2tensor.trax import backend\n",
        "from tensor2tensor.trax import layers as tl"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "height": 68
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 434,
          "status": "ok",
          "timestamp": 1563927323194,
          "user": {
            "displayName": "Lukasz Kaiser",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AAuE7mC8pChl87HbK_eOtVhtNPwUVx8btvfyYzH9UHn3=s64",
            "userId": "13267693649565518272"
          },
          "user_tz": 420
        },
        "id": "V09viOSEQvQe",
        "outputId": "822ec44b-1e2e-4b6e-9fc7-9fc29f5d3783"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[-7 -6 -5 -4 -3]\n",
            " [-2 -1  0  1  2]\n",
            " [ 3  4  5  6  7]]\n",
            "[[0 0 0 0 0]\n",
            " [0 0 0 1 2]\n",
            " [3 4 5 6 7]]\n"
          ]
        }
      ],
      "source": [
        "x  = onp.arange(-7, 8).reshape(3, -1)\n",
        "rng = backend.random.get_prng(0)\n",
        "layer = tl.Relu()\n",
        "params = layer.initialize(x.shape, x.dtype, rng)\n",
        "output = layer(x, params, rng=rng)\n",
        "print(x)\n",
        "print(output)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/dm_python:dm_notebook",
        "kind": "private"
      },
      "name": "Trax Layers",
      "provenance": [
        {
          "file_id": "1EH76AWQ_pvT4i8ZXfkv-SCV4MrmllEl5",
          "timestamp": 1563927451951
        }
      ],
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 2",
      "name": "python2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
